#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pymysql
from sshtunnel import SSHTunnelForwarder
from config import SSH_HOST, SSH_PORT, SSH_USERNAME, SSH_PASSWORD, DB_HOST, DB_PORT, DB_USERNAME, DB_PASSWORD

def create_ssh_tunnel():
    """
    使用sshtunnel库创建SSH隧道，返回SSHTunnelForwarder对象和本地端口
    """
    try:
        tunnel = SSHTunnelForwarder(
            (SSH_HOST, SSH_PORT),
            ssh_username=SSH_USERNAME,
            ssh_password=SSH_PASSWORD,
            remote_bind_address=(DB_HOST, DB_PORT),
            local_bind_address=('127.0.0.1', 0)  # 自动分配本地端口
        )
        tunnel.start()
        local_port = tunnel.local_bind_port
        print(f"SSH隧道已建立，本地端口: {local_port}")
        return tunnel, local_port
    except Exception as e:
        print(f"创建SSH隧道失败: {e}")
        return None, None

def test_database_connection():
    """
    测试数据库连接
    """
    tunnel, local_port = create_ssh_tunnel()
    if tunnel is None:
        print("Failed to create SSH tunnel. Exiting.")
        return

    database = 'tll_bi_dw'
    try:
        conn = pymysql.connect(
            host='127.0.0.1',
            port=local_port,
            user=DB_USERNAME,
            password=DB_PASSWORD,
            database=database
        )
        print("Database connection established successfully!")
        conn.close()
        print("Database connection closed.")
    except Exception as e:
        print(f"Database connection failed: {e}")
    finally:
        tunnel.stop()
        print("SSH tunnel closed.")

if __name__ == "__main__":
    test_database_connection()